#ifndef INPUT_HPP
#define INPUT_HPP

#include "Multigrid_One_Dimension.hpp"
#include "Multigrid_Two_Dimension.hpp"
#include "Function.hpp"
#include "json.hpp"
#include "Output.hpp"

using json = nlohmann::json;
void setMGParams_One_Dimension(MGParams_One_Dimension &mg_params, const json &input_data)
{
    try
    {
        int k = mg_params.grid_resolution;
        while (k > 2)
        {
            k = k / 2;
        }
        if (k == 1)
        {
            throw std::runtime_error("The grid_resolution must be '2^n'!!!");
        }

        mg_params.grid_resolution = mg_params.grid_resolution - 1;

        std::string restriction_str = input_data["MultiGrid_Params"]["restriction"].get<std::string>();
        if (restriction_str == "full_weighting")
        {
            mg_params.restriction = MGParams_One_Dimension::full_weighting;
        }
        else if (restriction_str == "injection")
        {
            mg_params.restriction = MGParams_One_Dimension::injection;
        }
        else
        {
            throw std::runtime_error("Restriction must be 'full_weighting' or 'injection'");
        }

        std::string interpolation_str = input_data["MultiGrid_Params"]["interpolation"].get<std::string>();
        if (interpolation_str == "linear")
        {
            mg_params.interpolation = MGParams_One_Dimension::linear;
        }
        else if (interpolation_str == "quadratic")
        {
            mg_params.interpolation = MGParams_One_Dimension::quadratic;
        }
        else
        {
            throw std::runtime_error("Interpolation must be 'linear' or 'quadratic'");
        }

        mg_params.v1 = input_data["MultiGrid_Params"]["v1"].get<int>();
        mg_params.v2 = input_data["MultiGrid_Params"]["v2"].get<int>();

        std::string type_str = input_data["MultiGrid_Params"]["type"].get<std::string>();
        if (type_str == "V_Cycles")
        {
            mg_params.type = MGParams_One_Dimension::V_Cycles;
        }
        else if (type_str == "FMG")
        {
            mg_params.type = MGParams_One_Dimension::FMG;
        }
        else
        {
            throw std::runtime_error("Type must be 'V_Cycles' or 'FMG'");
        }

        mg_params.max_iterations = input_data["MultiGrid_Params"]["max_iterations"].get<int>();
        mg_params.relative_accuracy = input_data["MultiGrid_Params"]["relative_accuracy"].get<double>();

        mg_params.recursions = input_data["MultiGrid_Params"]["recursion_limit"].get<int>();
        std::string initial_guess_str = input_data["MultiGrid_Params"]["initial_guess"].get<std::string>();
        if (initial_guess_str == "0")
        {
            mg_params.initial_guess = std::make_shared<InitialGuess0>();
        }
        else if (initial_guess_str == "1")
        {
            mg_params.initial_guess = std::make_shared<InitialGuess1>();
        }
        else
        {
            throw std::runtime_error("Initial guess must be '0' or '1'");
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << "Error in MultiGrid parameters: " << e.what() << std::endl;
        std::abort();
    }
}

void setBoundaryConditions_One_Dimension(BoundaryConditions_One_Dimension &bc, const json &input_data)
{
    try
    {
        std::string type_str = input_data["BoundaryConditions"]["type"].get<std::string>();
        if (type_str == "DIRICHLET")
        {
            bc.type = BoundaryConditions_One_Dimension::DIRICHLET;
        }
        else if (type_str == "NEUMANN")
        {
            bc.type = BoundaryConditions_One_Dimension::NEUMANN;
        }
        else if (type_str == "Mixed")
        {
            bc.type = BoundaryConditions_One_Dimension::Mixed;
        }
        else
        {
            throw std::runtime_error("Boundary condition type must be 'DIRICHLET', 'NEUMANN' or 'Mixed'");
        }

        bc.value = input_data["BoundaryConditions"]["value"].get<std::vector<double>>();

        if (bc.type == BoundaryConditions_One_Dimension::NEUMANN)
        {
            if (!input_data["BoundaryConditions"].contains("one_point"))
            {
                throw std::runtime_error("Missing 'one_point' value for Neumann boundary condition");
            }
            bc.one_point = input_data["BoundaryConditions"]["one_point"].get<std::vector<double>>();
        }

        if (bc.type == BoundaryConditions_One_Dimension::Mixed)
        {
            if (!input_data["BoundaryConditions"].contains("flag"))
            {
                throw std::runtime_error("Missing 'flag' for Mixed boundary conditions");
            }
            bc.flag = input_data["BoundaryConditions"]["flag"].get<std::vector<int>>();
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << "Error in Boundary Conditions: " << e.what() << std::endl;
        std::abort();
    }
}

void setsource_term_One_Dimension(std::shared_ptr<Function_One_Dimension> &source_term, const json &input_data)
{
    try
    {
        std::string source_term_str = input_data["source_term"].get<std::string>();
        if (source_term_str == "PI^2 * sin(PI x)")
        {
            source_term = std::make_shared<SourceTerm1>();
        }
        else if (source_term_str == "e^(x^2)")
        {
            source_term = std::make_shared<SourceTerm2>();
        }
        else if (source_term_str == "-(-sin(x) + cos^2(x) + 1) * exp(x + sin(x))")
        {
            source_term = std::make_shared<SourceTerm3>();
        }
        else if (source_term_str == "2 * PI^2 * sin(PI * x) * sin(PI * x)") {
            source_term = std::make_shared<SourceTerm4>();
        }
        else if (source_term_str == "-(-cos(x) - sin(x) + sin^2(x) + cos^2(x)) * exp(sin(x) + cos(x))") {
            source_term = std::make_shared<SourceTerm5>();
        }
        else
        {
            throw std::runtime_error("Source_term only supports 'PI^2 * sin(PI x)' , 2 * PI^2 * sin(PI * x) * sin(PI * x) and -(-sin(x) + cos^2(x) + 1) * exp(x + sin(x)) and -(-cos(x) - sin(x) + sin^2(x) + cos^2(x)) * exp(sin(x) + cos(x))!!!");
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << "Error in Boundary Conditions: " << e.what() << std::endl;
        std::abort();
    }
}

void setMGParams_Two_Dimension(MGParams_Two_Dimension &mg_params, const json &input_data)
{
    try
    {
        int k = mg_params.grid_resolution;
        while (k > 2)
        {
            k = k / 2;
        }
        if (k == 1)
        {
            throw std::runtime_error("The grid_resolution must be '2^n'!!!");
        }

        mg_params.grid_resolution = mg_params.grid_resolution - 1;

        std::string restriction_str = input_data["MultiGrid_Params"]["restriction"].get<std::string>();
        if (restriction_str == "full_weighting")
        {
            mg_params.restriction = MGParams_Two_Dimension::full_weighting;
        }
        else if (restriction_str == "injection")
        {
            mg_params.restriction = MGParams_Two_Dimension::injection;
        }
        else
        {
            throw std::runtime_error("Restriction must be 'full_weighting' or 'injection'");
        }

        std::string interpolation_str = input_data["MultiGrid_Params"]["interpolation"].get<std::string>();
        if (interpolation_str == "linear")
        {
            mg_params.interpolation = MGParams_Two_Dimension::linear;
        }
        else if (interpolation_str == "quadratic")
        {
            mg_params.interpolation = MGParams_Two_Dimension::quadratic;
        }
        else
        {
            throw std::runtime_error("Interpolation must be 'linear' or 'quadratic'");
        }

        mg_params.v1 = input_data["MultiGrid_Params"]["v1"].get<int>();
        mg_params.v2 = input_data["MultiGrid_Params"]["v2"].get<int>();

        std::string type_str = input_data["MultiGrid_Params"]["type"].get<std::string>();
        if (type_str == "V_Cycles")
        {
            mg_params.type = MGParams_Two_Dimension::V_Cycles;
        }
        else if (type_str == "FMG")
        {
            mg_params.type = MGParams_Two_Dimension::FMG;
        }
        else
        {
            throw std::runtime_error("Type must be 'V_Cycles' or 'FMG'");
        }

        mg_params.max_iterations = input_data["MultiGrid_Params"]["max_iterations"].get<int>();
        mg_params.relative_accuracy = input_data["MultiGrid_Params"]["relative_accuracy"].get<double>();
        mg_params.recursion_limit = input_data["MultiGrid_Params"]["recursion_limit"].get<int>();
        std::string initial_guess_str = input_data["MultiGrid_Params"]["initial_guess"].get<std::string>();
        if (initial_guess_str == "0")
        {
            mg_params.initial_guess = std::make_shared<InitialGuess0_2d>();
        }
        else if (initial_guess_str == "1")
        {
            mg_params.initial_guess = std::make_shared<InitialGuess1_2d>();
        }
        else
        {
            throw std::runtime_error("Initial guess must be '0' or '1'");
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << "Error in MultiGrid parameters: " << e.what() << std::endl;
        std::abort();
    }
}

void setBoundaryConditions_Two_Dimension(BoundaryConditions_Two_Dimension &bc, const json &input_data)
{
    try
    {
        std::string type_str = input_data["BoundaryConditions"]["type"].get<std::string>();
        if (input_data["BoundaryConditions"]["u_1"] == "exp(0 + sin(x))")
        {
            bc.u_1 = std::make_shared<u1_1_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_1"] == "exp(0 + sin(x))")
        {
            bc.u_1 = std::make_shared<u1_partial_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_1"] == "0")
        {
            bc.u_1 = std::make_shared<u1_2_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_1"] == "PI * sin(PI * x)")
        {
            bc.u_1 = std::make_shared<u1_2_partial_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_1"] == "exp(sin(x) + 1)")
        {
            bc.u_1 = std::make_shared<u1_3_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_1"] == "0")
        {
            bc.u_1 = std::make_shared<u1_3_partial_2d>();
        }
        else
        {
            throw std::runtime_error("For BoundaryConditions_Two_Dimension, u_1 only applies 'exp(0 + sin(x))' , '0' , 'PI * sin(PI * x)' and 'exp(sin(x) + 1)'");
        }

        if (input_data["BoundaryConditions"]["u_2"] == "exp(y + sin(1))")
        {
            bc.u_2 = std::make_shared<u2_1_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_2"] == "cos(1) * exp(y + sin(1))")
        {
            bc.u_2 = std::make_shared<u2_partial_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_2"] == "0")
        {
            bc.u_2 = std::make_shared<u2_2_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_2"] == "-PI * sin(PI * y)")
        {
            bc.u_2 = std::make_shared<u2_2_partial_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_2"] == "exp(sin(1) + cos(y))")
        {
            bc.u_2 = std::make_shared<u2_3_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_2"] == "exp(sin(1) + cos(y)) * cos(1)")
        {
            bc.u_2 = std::make_shared<u2_3_partial_2d>();
        }
        else
        {
            throw std::runtime_error("For BoundaryConditions_Two_Dimension, u_2 only applies 'exp(y + sin(1))' , 'cos(1) * exp(y + sin(1))', '0' , '-PI * sin(PI * y)' , 'exp(sin(1) + cos(y))' and 'exp(sin(1) + cos(y)) * cos(1)'");
        }

        if (input_data["BoundaryConditions"]["u_3"] == "exp(1 + sin(x))")
        {
            bc.u_3 = std::make_shared<u3_1_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_3"] == "exp(1 + sin(x))")
        {
            bc.u_3 = std::make_shared<u3_partial_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_3"] == "0")
        {
            bc.u_3 = std::make_shared<u3_2_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_3"] == "-PI * sin(PI * x)")
        {
            bc.u_3 = std::make_shared<u3_2_partial_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_3"] == "exp(sin(x) + cos(1))")
        {
            bc.u_3 = std::make_shared<u3_3_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_3"] == "exp(sin(x) + cos(1)) * (-sin(1))")
        {
            bc.u_3 = std::make_shared<u3_3_partial_2d>();
        }
        else
        {
            throw std::runtime_error("For BoundaryConditions_Two_Dimension, u_3 only applies 'exp(1 + sin(x))' , 'exp(1 + sin(x))' , '0' , '-PI * sin(PI * x)' , 'exp(sin(x) + cos(1))' and 'exp(sin(x) + cos(1)) * (-sin(1))'");
        }

        // u_4
        if (input_data["BoundaryConditions"]["u_4"] == "exp(y + sin(0))")
        {
            bc.u_4 = std::make_shared<u4_1_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_4"] == "cos(0) * exp(y + sin(0))")
        {
            bc.u_4 = std::make_shared<u4_partial_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_4"] == "0")
        {
            bc.u_4 = std::make_shared<u4_2_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_4"] == "PI * sin(PI * y)")
        {
            bc.u_4 = std::make_shared<u4_2_partial_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_4"] == "exp(cos(y))")
        {
            bc.u_4 = std::make_shared<u4_3_2d>();
        }
        else if (input_data["BoundaryConditions"]["u_4"] == "exp(cos(y))")
        {
            bc.u_4 = std::make_shared<u4_3_partial_2d>();
        }
        else
        {
            throw std::runtime_error("For BoundaryConditions_Two_Dimension, u_4 only applies 'exp(y + sin(0))' , 'cos(0) * exp(y + sin(0))' , '0' , 'PI * sin(PI * y)' and 'exp(cos(y))'");
        }

        if (type_str == "DIRICHLET")
        {
            bc.type = BoundaryConditions_Two_Dimension::DIRICHLET;
        }
        else if (type_str == "NEUMANN")
        {
            bc.type = BoundaryConditions_Two_Dimension::NEUMANN;
        }
        else if (type_str == "Mixed")
        {
            bc.type = BoundaryConditions_Two_Dimension::Mixed;
        }
        else
        {
            throw std::runtime_error("Boundary condition type must be 'DIRICHLET', 'NEUMANN' or 'Mixed'");
        }

        if (bc.type == BoundaryConditions_Two_Dimension::NEUMANN)
        {
            if (!input_data["BoundaryConditions"].contains("one_point"))
            {
                throw std::runtime_error("Missing 'one_point' value for Neumann boundary condition");
            }
            bc.one_point = input_data["BoundaryConditions"]["one_point"].get<std::vector<double>>();
        }

        if (bc.type == BoundaryConditions_Two_Dimension::Mixed)
        {
            if (!input_data["BoundaryConditions"].contains("flag"))
            {
                throw std::runtime_error("Missing 'flag' for Mixed boundary conditions");
            }
            bc.flag = input_data["BoundaryConditions"]["flag"].get<std::vector<int>>();
            if (bc.flag.size() != 4)
            {
                throw std::runtime_error("The size of flag  for Mixed boundary condition must be 4");
            }
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << "Error in Boundary Conditions: " << e.what() << std::endl;
        std::abort();
    }
}

void setsource_term_Two_Dimension(std::shared_ptr<Function_Two_Dimension> &source_term, const json &input_data)
{
    try
    {
        std::string source_term_str = input_data["source_term"].get<std::string>();
        if (source_term_str == "-(-sin(x) + cos^2(x) + 1) * exp(y + sin(x))")
        {
            source_term = std::make_shared<SourceTerm_1_2d>();
        }
        else if (source_term_str == "2 * PI^2 * sin(PI * x) * sin(PI * y)")
        {
            source_term = std::make_shared<SourceTerm_2_2d>();
        }
        else if (source_term_str == "-(-cos(y) - sin(x) + sin^2(y) + cos^2(x)) * exp(sin(x) + cos(y))")
        {
            source_term = std::make_shared<SourceTerm_3_2d>();
        }
        else
        {
            throw std::runtime_error("Source_term only supports '-(-sin(x) + cos^2(x) + 1) * exp(y + sin(x))' , '2 * PI^2 * sin(PI * x) * sin(PI * y)' and '-(-cos(y) - sin(x) + sin^2(y) + cos^2(x)) * exp(sin(x) + cos(y))'!!!");
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << "Error in Boundary Conditions: " << e.what() << std::endl;
        std::abort();
    }
}

void Initialization()
{
    // Read input file
    std::ifstream input_file("../Input.json");
    if (!input_file.is_open())
    {
        std::cerr << "Error: Could not open input file!" << std::endl;
        std::abort();
    }

    json input_data;
    input_file >> input_data;
    int dimension = input_data["Dimension"].get<int>();

    if (dimension == 1)
    {
        double error = 0;
        double temp = 0;
        std::vector<int> k = input_data["MultiGrid_Params"]["grid_resolution"].get<std::vector<int>>();
        for (size_t i = 0; i < k.size(); i++)
        {
            MGParams_One_Dimension mg_params;
            mg_params.grid_resolution = k[i];
            BoundaryConditions_One_Dimension bc;
            std::shared_ptr<Function_One_Dimension> source_term;
            setMGParams_One_Dimension(mg_params, input_data);
            setBoundaryConditions_One_Dimension(bc, input_data);
            setsource_term_One_Dimension(source_term, input_data);
            Multigrid_One_Dimension multigrid(mg_params, bc, source_term);
            multigrid.solve();
            const Eigen::VectorXd &solution = multigrid.GetSolution();
            if (i == 0)
            {
                temp = output_1d(solution);
                error = temp;
                std::cout << "L2 Error: " << error << std::endl;
            }
            else
            {
                temp = output_1d(solution);
                error /= temp;
                std::cout << "L2 Error: " << temp << std::endl;
                std::cout << "Error convergence rates: " << error << std::endl;
            }
            error = temp;
        }
    }
    else if (dimension == 2)
    {
        double error = 0;
        double temp = 0;
        std::vector<int> k = input_data["MultiGrid_Params"]["grid_resolution"].get<std::vector<int>>();
        for (size_t i = 0; i < k.size(); i++)
        {
            MGParams_Two_Dimension mg_params;
            mg_params.grid_resolution = k[i];
            BoundaryConditions_Two_Dimension bc;
            std::shared_ptr<Function_Two_Dimension> source_term;
            setMGParams_Two_Dimension(mg_params, input_data);
            setsource_term_Two_Dimension(source_term, input_data);
            setBoundaryConditions_Two_Dimension(bc, input_data);
            Multigrid_Two_Dimension multigrid(mg_params, bc, source_term);
            multigrid.solve();
            const Eigen::VectorXd &solution = multigrid.GetSolution();
            if (i == 0)
            {
                temp = output_2d(solution);
                std::cout << "L2 Error: " << temp << std::endl;
            }
            else
            {
                temp = output_2d(solution);
                error /= temp;
                std::cout << "L2 Error: " << temp << std::endl;
                std::cout << "Error convergence rates: " << error << std::endl;
            }
            error = temp;
        }
    }
    else
    {
        std::cerr << "Error: The project only supports one and two dimensions!" << std::endl;
        std::abort();
    }
}

#endif // INPUT_HPP